import os
import baostock as bs
import tushare as ts
import pandas as pd
import numpy as np
from tqdm import tqdm
import time

# 定义数据的时间范围
data_start_date = '2021-01-01'  # 数据开始日期
data_end_date = '2024-06-02'  # 数据结束日期
train_data_start_date = '2018-01-01'  # 训练数据开始日期
train_data_end_date = '2021-12-31'  # 训练数据结束日期

# 定义存储路径
base_data_path = './data/'  # 基础数据存储路径
data_path = './data/stocks/'  # 股票数据存储路径
train_data_path = './data/train_data/'  # 训练数据存储路径

# 使用tushare数据API
apikey = ''  # Tushare API Key
if os.path.exists('tushare_apikey.txt'):  # 如果存在API Key文件，则读取API Key
    with open('tushare_apikey.txt', 'r') as f:
        apikey = f.read().strip()  # 读取API Key并去除首尾空白字符
pro = ts.pro_api(apikey)  # 创建Tushare Pro API接口对象

# 设置是否需要下载数据
download_stocks = True  # 是否下载个股数据
download_indexes = True  # 是否下载指数数据

# 定义下载个股所需字段
fields = "date,code,open,high,low,close,preclose,volume,amount,turn,peTTM,psTTM,pcfNcfTTM,pbMRQ"
ts_basic_fields = "trade_date,total_share,float_share,free_share,total_mv,circ_mv"  # Tushare基本指标字段

# 登录Baostock
lg = bs.login()
# 检查登录状态
if lg.error_code != '0':  # 如果登录失败
    print('login respond error_code:' + lg.error_code)
    print('login respond  error_msg:' + lg.error_msg)


# 定义检查Baostock请求状态的函数
def status(raw):
    if raw.error_code != '0':  # 如果请求失败
        print('query_history_k_data_plus respond error_code:' + raw.error_code)
        print('query_history_k_data_plus respond  error_msg:' + raw.error_msg)


# 定义从Baostock加载数据的函数
def data_load(raw):
    """从Baostock响应中读取数据并转换成DataFrame"""
    data = []
    while (raw.error_code == '0') & raw.next():  # 当没有错误且有下一条记录时
        data.append(raw.get_row_data())  # 添加一条记录到列表
    df = pd.DataFrame(data, columns=raw.fields)  # 将列表转换为DataFrame
    return df


# 定义将Baostock股票代码转换为Tushare股票代码的函数
def ts_c(bs_stock_code: str) -> str:
    """将Baostock股票代码转换为Tushare股票代码格式"""
    return bs_stock_code[3:] + '.' + bs_stock_code[0:2].upper()


# 定义将Baostock日期转换为Tushare日期的函数
def ts_d(date):
    """将Baostock日期格式（YYYY-MM-DD）转换为Tushare日期格式（YYYYMMDD）"""
    return date.replace('-', '')


# 定义将Tushare日期序列转换为Baostock日期序列的函数
def bs_d(ts_date_series):
    """将Tushare日期序列转换为Baostock日期序列"""
    date_series = []
    for ts_date in ts_date_series:
        date = str(ts_date)
        date_series.append(date[:4] + '-' + date[4:6] + '-' + date[6:])  # 转换日期格式
    return date_series


# 如果设置了下载个股数据
if download_stocks:
    rs = bs.query_hs300_stocks()  # 查询沪深300成分股
    # 获取股票名称与代码，并保存至CSV文件
    result = data_load(rs)
    result.to_csv(base_data_path + 'hs300_stocks.csv')

    # 对于每一只股票，下载其历史K线数据和基本指标数据
    for code in tqdm(result['code']):
        for start_date, end_date, path in ((data_start_date, data_end_date, data_path),
                                           (train_data_start_date, train_data_end_date, train_data_path)):
            # 查询历史K线数据
            rs = bs.query_history_k_data_plus(code, fields,
                                              start_date=start_date, end_date=end_date,
                                              frequency="d", adjustflag="3")
            status(rs)  # 检查请求状态

            # 加载Baostock日线数据
            df1 = data_load(rs)
            df1.set_index('date', inplace=True)  # 将日期设置为索引

            # 查询Tushare的日线基本指标数据
            df2 = pro.daily_basic(ts_code=ts_c(code), start_date=ts_d(start_date),
                                  end_date=ts_d(end_date), fields=ts_basic_fields)
            df2['trade_date'] = bs_d(df2['trade_date'])  # 转换日期格式
            df2.set_index('trade_date', inplace=True)  # 将交易日期设置为索引

            # 合并两个DataFrame
            df1 = df1.join(df2)

            # 处理缺失值
            df1.replace(to_replace=r'^\s*$', value=np.nan, regex=True, inplace=True)  # 将空白字符串替换为NaN
            df1.fillna(0, inplace=True)  # 用0填充NaN值

            # 保存数据到CSV文件
            df1.to_csv(path + code + '.csv')
            del df1, df2  # 删除变量以释放内存

        time.sleep(0.2)  # 避免Tushare API调用频率限制

# 如果设置了下载指数数据
if download_indexes:
    # 查询沪深300指数的历史K线数据
    rs = bs.query_history_k_data_plus('sh.000300',
                                      "date,code,open,high,low,close,preclose,volume,amount,pctChg",
                                      start_date=data_start_date, end_date=data_end_date, frequency="d")
    status(rs)  # 检查请求状态
    # 加载数据并保存到CSV文件
    data_load(rs).to_csv(data_path + 'sh.000300.csv')

# 登出Baostock
bs.logout()